"""
角色链代理模块
"""

import os
import logging
from typing import Dict, Any, List, Optional, AsyncGenerator, Union
import json

from api.models.knowledge_base import KnowledgeBase
from api.models.llm_client import LLMClient
from api.models.web_search import WebSearchTool

# 设置日志
logger = logging.getLogger("RoleChainAgent")
logger.setLevel(logging.INFO)

class RoleChainAgent:
    """基于LangChain的角色扮演Agent"""
    
    def __init__(self, knowledge_dir="knowledge"):
        """初始化Agent
        
        Args:
            knowledge_dir: 知识库目录
        """
        self.logger = logging.getLogger("RoleChainAgent")
        self.logger.setLevel(logging.INFO)
        
        # 初始化知识库
        self.knowledge_base = KnowledgeBase(knowledge_dir=knowledge_dir)
        

        # 初始化网络搜索工具
        self.web_search_tool = WebSearchTool()
        
        # 初始化LLM客户端
        self.llm_client = LLMClient()
        
        # 当前角色信息
        self.current_role = None
        
        self.logger.info("RoleChainAgent初始化完成")
    
    async def get_knowledge_for_role(self, role_id: str, query: str) -> str:
        """根据角色名称获取知识库信息
        
        Args:
            role_id: 角色ID
            query: 用户查询
            
        Returns:
            知识库检索结果
        """
        # 将角色ID转换为目录名格式
        # role_id = role_id.lower().replace(" ", "_")
        
        # 检查知识库目录是否存在
        knowledge_dir = os.path.join(self.knowledge_base.knowledge_dir, role_id)
        if not os.path.exists(knowledge_dir) or not os.path.isdir(knowledge_dir):
            self.logger.info(f"角色 {role_id} 的知识库目录不存在，跳过检索")
            return ""
                
        # 搜索知识库
        search_results = self.knowledge_base.search(query, role_id=role_id)
        
        if not search_results:
            self.logger.info(f"未找到角色 {role_id} 的知识库信息")
            return ""
        
        # 合并搜索结果
        knowledge_text = "\n\n".join([result["content"] for result in search_results])
        self.logger.info(f"找到角色 {role_id} 的知识库信息: {len(search_results)} 条")
        
        return knowledge_text
        
    async def search_web(self, query: str, max_results: int = 3) -> str:
        """执行网络搜索
        
        Args:
            query: 搜索查询
            max_results: 最大结果数量
            
        Returns:
            搜索结果文本
        """
        try:
            # 使用百度AI搜索API
            async with self.web_search_tool as search_tool:
                search_results = await search_tool.search_baidu_ai(query, max_results=max_results)
                
                # 如果百度AI搜索失败，尝试普通百度搜索
                if not search_results:
                    self.logger.info("百度AI搜索未返回结果，尝试普通百度搜索")
                    search_results = await search_tool.search_baidu(query, max_results=max_results)
            
            if not search_results:
                self.logger.info(f"网络搜索未找到结果: {query}")
                return ""
            
            # 格式化搜索结果
            formatted_results = []
            for i, result in enumerate(search_results):
                title = result.get('title', '无标题')
                snippet = result.get('snippet', '无摘要')
                url = result.get('url', '无链接')
                formatted_results.append(f"[{i+1}] {title}\n摘要: {snippet}\n来源: {url}")
            
            web_search_text = "\n\n".join(formatted_results)
            self.logger.info(f"网络搜索找到 {len(search_results)} 条结果")
            
            return web_search_text
            
        except Exception as e:
            self.logger.error(f"网络搜索出错: {str(e)}")
            return ""
    
    async def process_custom_role(self, role_info: Dict[str, Any], user_input: str, history: List[Dict[str, str]] = None, websh: bool = True):
        """处理自定义角色请求
        
        Args:
            role_info: 角色信息，包含roleId, roleName, systemPrompt等
            user_input: 用户输入
            history: 历史对话
            websh: 是否使用网络搜索，默认为True
            
        Returns:
            生成的回复
        """
        try:
            # 确保role_info是字典类型
            if not isinstance(role_info, dict):
                role_info = role_info.__dict__ if hasattr(role_info, '__dict__') else {"roleId": "", "roleName": "", "systemPrompt": ""}
                
            role_id = role_info.get("roleId", "")
            role_name = role_info.get("roleName", "")
            system_prompt = role_info.get("systemPrompt", "")
            
            # 获取知识库信息
            knowledge_text = await self.get_knowledge_for_role(role_id, user_input)
            
            # 如果有知识库信息，将其添加到用户输入中
            enhanced_input = user_input
            
            # 执行网络搜索（如果websh为True）
            web_search_text = ""
            if websh:
                web_search_text = await self.search_web(user_input)
                self.logger.info(f"网络搜索已启用，获取到结果: {len(web_search_text) > 0}")
            else:
                self.logger.info("网络搜索已禁用")
            
            # 组合知识库信息和网络搜索结果
            if knowledge_text and web_search_text:
                enhanced_input = f"以下是与问题相关的知识库信息:\n{knowledge_text}\n\n以下是网络搜索结果:\n{web_search_text}\n\n用户问题: {user_input}"
                self.logger.info(f"增强后的用户输入(知识库+网络搜索): {enhanced_input[:100]}...")
            elif knowledge_text:
                enhanced_input = f"以下是与问题相关的知识库信息:\n{knowledge_text}\n\n用户问题: {user_input}"
                self.logger.info(f"增强后的用户输入(仅知识库): {enhanced_input[:100]}...")
            elif web_search_text:
                enhanced_input = f"以下是网络搜索结果:\n{web_search_text}\n\n用户问题: {user_input}"
                self.logger.info(f"增强后的用户输入(仅网络搜索): {enhanced_input[:100]}...")
            
            # 准备消息历史
            messages = []
            if history:
                for msg in history:
                    # 获取角色和内容（确保键存在，避免KeyError）
                    role = msg.get("role")
                    content = msg.get("content", "")
                    if role == "user":
                        # 用户消息
                        messages.append({"role": "user", "content": content})
                    elif role == role_id:
                        # 助手消息（匹配角色ID）
                        messages.append({"role": "assistant", "content": content})
            
            # 添加当前用户消息
            messages.append({"role": "user", "content": enhanced_input})
            

            
            # 初始化LLM客户端
            llm_client = LLMClient()
            
            # 调用LLM客户端的流式生成方法
            async for chunk in llm_client.generate_text_stream(
                messages=messages,
                system_prompt=system_prompt
            ):
                yield chunk
                
        except Exception as e:
            self.logger.error(f"处理自定义角色请求出错: {str(e)}")
            yield f"处理请求时出错: {str(e)}"